--[[
Ip2region lua binding

@author chenxin<chenxin619315@gmail.com>
@date   2018/10/02
]]--

local bit32 = require("bit");

local INDEX_BLOCK_LENGTH  = 12;
local TOTAL_HEADER_LENGTH = 8192;
local _M = {
  dbFile = "",
  dbFileHandler = "",
  dbBinStr = "",
  HeaderSip = "",
  HeaderPtr = "",
  headerLen = 0,
  firstIndexPtr = 0,
  lastIndexPtr = 0,
  totalBlocks = 0
};

_G["Ip2region"] = _M;

-- set the __index to itself
_M.__index = _M;

-- set the print meta-method
_M.__tostring = function(table)
  local t = {
    "dbFile=" .. table.dbFile,
    -- "dbFileHandler=" .. table.dbFileHandler,
    "headerLen=" .. table.headerLen,
    "firstIndexPtr" .. table.firstIndexPtr,
    "lastIndexPtr" .. table.lastIndexPtr,
    "totalBlocks" .. table.totalBlocks
  };

  return table.concat(t, ",");
end

--[[
construct method

@param  obj
@return Ip2region object
--]]
function _M.new(dbFile)
  obj = {};
  setmetatable(obj, _M);
  obj.dbFile = dbFile;
  return obj;
end

-- 将2个有符号32位值进行无符号方式比较
-- n1 > n2 时返回1，
-- n1 < n2 时返回-1
-- n1 == n2 时返回0
function compareLong(n1,n2)
  local h1,h2,l1,l2;
  h1 = bit32.rshift(n1,16);
  h2 = bit32.rshift(n2,16);
  if(h1 > h2 ) then
    return 1;
  elseif ( h1 < h2) then
    return -1;
  else
    l1 = bit32.band(n1,0xffff);
    l2 = bit32.band(n2,0xffff);
    if(l1 > l2 ) then
      return 1;
    elseif( l1 < l2) then
      return -1;
    else
      return 0;
    end
  end
end
--[[
internal function to get a integer from a binary string

@param  dbBinStr
@param  idx
@return Integer
]]--
function getLong(bs, idx)
  local bit = string.byte(string.sub(bs, idx, idx));
  if( bit == nil) then
    bit = 0;
  end
  local a1 = bit;
  bit = string.byte(string.sub(bs, idx+1, idx+1));
  if ( bit == nil) then
    bit = 0;
  end
  local a2 = bit32.lshift(bit,  8);
  bit = string.byte(string.sub(bs, idx+2, idx+2));
  if(bit == nil) then
    bit = 0;
  end
  local a3 = bit32.lshift(bit, 16);
  bit = string.byte(string.sub(bs, idx+3, idx+3));
  if(bit == nil) then
    bit = 0;
  end
  local a4 = bit32.lshift(bit, 24);

  local val = bit32.bor(a1, a2);
  val = bit32.bor(val, a3);
  val = bit32.bor(val, a4);

  return val;
end


--[[
internal function to convert the string ip to a long value

@param  ip
@return Integer
]]--
function _M.ip2long(self, ip)
  -- dynamic arguments checking
  -- to support object.ip2long and object:ip2long access
  if ( type(self) == "string") then
    ip = self;
  end

  local ini = 1;
  local iip = 0;
  local off = 24;
  while true do
    local pos = string.find(ip, '.', ini, true);
    if ( not pos ) then
      break;
    end

    local sub = string.sub(ip, ini, pos - 1);
    if ( string.len(sub) < 1 ) then
      return nil;
    end

    iip = bit32.bor(iip, bit32.lshift(tonumber(sub), off));
    ini = pos + 1;
    off = off - 8;
  end

  -- check if it is a valid ip address
  if ( off ~= 0 or ini > string.len(ip) ) then
    return nil;
  end

  local sub = string.sub(ip, ini);
  if ( string.len(sub) < 1 ) then
    return nil;
  end

  return bit32.bor(iip, bit32.lshift(tonumber(sub), off));
end


--[[
internal function to get the whole content of a file

@param  file
@return String
]]--
function get_file_contents(file)
  local fi = io.input(file);
  if ( not fi ) then
    return nil;
  end

  local str = io.read("*a");
  io.close();
  return str;
end


--[[
set the current db file path

@param  dbFile
]]--
function _M:setDbFile(dbFile)
  self.dbFile = dbFile;
end


--[[
all the db binary string will be loaded into memory
then search the memory only and this will a lot faster than disk base search
@Note: invoke it once before put it to public invoke could make it thread safe

@param  ip
@return table or nil for failed
]]--
function _M:memorySearch(ip)
  -- string ip conversion
  if ( type(ip) == "string" ) then
    ip = bit.tobit(self:ip2long(ip));
    if ( ip == nil ) then
      return nil;
    end
  end;

  -- check and load the binary string for the first time
  if ( self.dbBinStr == "" ) then
    self.dbBinStr = get_file_contents(self.dbFile);
    if ( not self.dbBinStr ) then
      return nil;
    end

    self.firstIndexPtr = getLong(self.dbBinStr, 1);
    self.lastIndexPtr  = getLong(self.dbBinStr, 5);
    self.totalBlocks   = (self.lastIndexPtr - self.firstIndexPtr)/INDEX_BLOCK_LENGTH + 1;
  end

  --ngx.say('firstIndexPtr=',self.firstIndexPtr);

  -- binary search to define the data
  local l = 0;
  local h = self.totalBlocks;
  local dataPtr = 0;
  ngx.say('ip:',ip);
  while ( l <= h ) do
    local m = math.floor((l + h) / 2);
    local p = self.firstIndexPtr + m * INDEX_BLOCK_LENGTH;
    ngx.say('p=',p);
    local sip = getLong(self.dbBinStr, p + 1);
    ngx.say('sip=',sip);
    if ( compareLong(ip,sip) < 0 ) then
      h = m - 1;
    else
      local eip = getLong(self.dbBinStr, p + 5);  -- 4 + 1
      ngx.say('eip=',eip);
      if ( compareLong(ip,eip) > 0 ) then
        l = m + 1;
        ngx.say('l=',l);
        ngx.say('h=',h);
      else
        dataPtr = getLong(self.dbBinStr, p + 9); -- 8 + 1
        break;
      end
    end
  end

  -- not matched just stop it here
  if ( dataPtr == 0 ) then return nil end

  -- get the data
  local dataLen = bit32.band(bit32.rshift(dataPtr, 24), 0xFF);
  dataPtr = bit32.band(dataPtr, 0x00FFFFFF);
  local dptr = dataPtr + 5;   -- 4 + 1

  return {
    city_id = getLong(self.dbBinStr, dataPtr),
    region  = string.sub(self.dbBinStr, dptr, dptr + dataLen - 5)
  };
end


--[[
get the data block through the specified ip address
or long ip numeric with binary search algorithm

@param  ip
@return table or nil for failed
]]--
function _M:binarySearch(ip)
  -- check and conver the ip address
  if ( type(ip) == "string" ) then
    ip = self:ip2long(ip);
    if ( ip == nil ) then
      return nil;
    end
  end

  if ( self.totalBlocks == 0 ) then
    -- check and open the original db file
    self.dbFileHandler = io.open(self.dbFile, "r");
    if ( not self.dbFileHandler ) then
      return nil;
    end

    self.dbFileHandler:seek("set", 0);
    local superBlock = self.dbFileHandler:read(8);

    self.firstIndexPtr = getLong(superBlock, 1);    -- 0 + 1
    self.lastIndexPtr  = getLong(superBlock, 5);    -- 4 + 1
    self.totalBlocks   = (self.lastIndexPtr-self.firstIndexPtr)/INDEX_BLOCK_LENGTH + 1;
  end


  -- binary search to define the data
  local l = 0;
  local h = self.totalBlocks;
  local dataPtr = 0;
  while ( l <= h ) do
    local m = math.floor((l + h) / 2);
    local p = m * INDEX_BLOCK_LENGTH;
    self.dbFileHandler:seek("set", self.firstIndexPtr + p);
    local buffer = self.dbFileHandler:read(INDEX_BLOCK_LENGTH);
    local sip = getLong(buffer, 1); -- 0 + 1
    if ( ip < sip ) then
      h = m - 1;
    else
      local eip = getLong(buffer, 5);   -- 4 + 1
      if ( ip > eip ) then
        l = m + 1;
      else
        dataPtr = getLong(buffer, 9);   -- 8 + 1
        break;
      end
    end
  end

  -- not matched just stop it here
  if ( dataPtr == 0 ) then return nil; end


  -- get the data
  local dataLen = bit32.band(bit32.rshift(dataPtr, 24), 0xFF);
  dataPtr = bit32.band(dataPtr, 0x00FFFFFF);

  self.dbFileHandler:seek("set", dataPtr);
  local data = self.dbFileHandler:read(dataLen);

  return {
    city_id = getLong(data, 1),    -- 0 + 1
    region  = string.sub(data, 5)  -- 4 + 1
  };
end


--[[
get the data block associated with the specified ip with b-tree search algorithm

@param  ip
@return table or nil for failed
]]--
function _M:btreeSearch(ip)
  -- string ip to integer conversion
  if ( type(ip) == "string" ) then
    ip = self:ip2long(ip);
    if ( ip == nil ) then
      return nil;
    end
  end

  -- check and load the header
  if ( self.headerLen == 0 ) then
    -- check and open the original db file
    self.dbFileHandler = io.open(self.dbFile, 'r');
    if ( not self.dbFileHandler ) then
      return nil;
    end

    self.dbFileHandler:seek("set", 8);
    local buffer = self.dbFileHandler:read(TOTAL_HEADER_LENGTH);

    -- fill the header
    local i = 0;
    local idx = 0;
    self.HeaderSip = {};
    self.HeaderPtr = {};
    for i=0, TOTAL_HEADER_LENGTH, 8 do
      local startIp = getLong(buffer, i + 1); -- 0 + 1
      local dataPtr = getLong(buffer, i + 5); -- 4 + 1
      if ( dataPtr == 0 ) then
        break;
      end

      table.insert(self.HeaderSip, startIp);
      table.insert(self.HeaderPtr, dataPtr);
      idx = idx + 1;
    end

    self.headerLen = idx;
  end


  -- 1. define the index block with the binary search
  local l = 0;
  local h = self.headerLen;
  local sptr = 0;
  local eptr = 0;
  while ( l <= h ) do
    local m = math.floor((l + h) / 2);
    -- perfetc matched, just return it
    if ( compareLong(ip,self.HeaderSip[m] ) == 0) then
      if ( m > 0 ) then
        sptr = self.HeaderPtr[m-1];
        eptr = self.HeaderPtr[m  ];
      else
        sptr = self.HeaderPtr[m  ];
        eptr = self.HeaderPtr[m+1];
      end

      break;
    end

    -- less then the middle value
    if ( compareLong(ip,self.HeaderSip[m]) < 0 ) then
      if ( m == 0 ) then
        sptr = self.HeaderPtr[m  ];
        eptr = self.HeaderPtr[m+1];
        break;
      elseif ( compareLong(ip , self.HeaderSip[m-1]) > 0 ) then
        sptr = self.HeaderPtr[m-1];
        eptr = self.HeaderPtr[m  ];
        break;
      end
      h = m - 1;
    else
      if ( m == self.headerLen - 1 ) then
        sptr = self.HeaderPtr[m-1];
        eptr = self.HeaderPtr[m  ];
        break;
      elseif ( compareLong (ip , self.HeaderSip[m+1]) <= 0 ) then
        sptr = self.HeaderPtr[m  ];
        eptr = self.HeaderPtr[m+1];
        break;
      end
      l = m + 1;
    end
  end

  -- match nothing just stop it
  if ( sptr == 0 ) then return nil; end

  -- 2. search the index blocks to define the data
  self.dbFileHandler:seek("set", sptr);
  local blockLen = eptr - sptr;
  local index = self.dbFileHandler:read(blockLen + INDEX_BLOCK_LENGTH);
  local dataPtr = 0;

  l = 0;
  h = blockLen / INDEX_BLOCK_LENGTH;
  while ( l <= h ) do
    local m = math.floor((l + h) / 2);
    local p = m * INDEX_BLOCK_LENGTH;
    local sip = getLong(index, p + 1);       -- 0 + 1
    if ( compareLong(ip,sip) < 0) then
      h = m - 1;
    else
      local eip = getLong(index, p + 5);   -- 4 + 1
      if ( compareLong(ip,eip) > 0 ) then
        l = m + 1;
      else
        dataPtr = getLong(index, p + 9); -- 8 + 1
        break;
      end
    end
  end

  -- not matched
  if ( dataPtr == 0 ) then return nil; end

  -- 3. get the data
  local dataLen = bit32.band(bit32.rshift(dataPtr, 24), 0xFF);
  dataPtr = bit32.band(dataPtr, 0x00FFFFFF);

  self.dbFileHandler:seek("set", dataPtr);
  local data = self.dbFileHandler:read(dataLen);

  return {
    city_id = getLong(data, 1),     -- 0 + 1
    region  = string.sub(data, 5)   -- 4 + 1
  };
end


--[[
close the object and do the basic gc
]]--
function _M.close(self)
  if ( self.dbFileHandler ~= "" ) then
    self.dbFileHandler:close();
  end

  if ( self.dbBinStr ~= "" ) then
    self.dbBinStr = nil;
  end
end


return _M;
